# -- coding:utf-8 --
file_path = "/Users/mac/Downloads/datasets/cmrc2018/"
import numpy as np
from torch.utils.data import DataLoader
import torch

# %%
from transformers import AutoTokenizer

checkpoint = 'bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# %%
from torch import nn
import torch
from transformers import AutoConfig
from transformers import BertPreTrainedModel, BertModel

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device')

class BertForExtractiveQA(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.bert = BertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        self.post_init()
    
    def forward(self, x):
        bert_output = self.bert(**x)
        sequence_output = bert_output.last_hidden_state
        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1).contiguous()
        end_logits = end_logits.squeeze(-1).contiguous()

        return start_logits, end_logits
    
config = AutoConfig.from_pretrained(checkpoint)
config.num_labels=2
model = BertForExtractiveQA.from_pretrained(checkpoint, config=config).to(device)
print(model)

import collections
from transformers import AdamW, get_scheduler
import os
import random
from tqdm.auto import tqdm
import json

model.load_state_dict(torch.load(file_path + 'epoch_2_valid_avg_74.3551_model_weights.bin', map_location=torch.device('cpu')))
model.eval()

# %%

# raw_str = '{"title": "现代军事包括什么", "text": "军事学与甚多范畴有关，主要与战争有关。此外，军事学本身包含了各种学问。军事是政治的一部分，战争是政治的一种延续，是一国或者集团用暴力手段达到自己目标和目的的方式，而目标和目的往往与利益有关。战争是军事的集中体现，但不是唯一的体现。第二次世界大战（1939年－1945年）后的美国和苏联冷战，就是一种威慑基础上的回避战争方式的斗争。在人类可以看到的未来，军事始终是政治生活中重要的方面，并在科学技术上对人类生活予重大影响：人类很多科技成就往往先产生于军事领域然后普及到非军事领域的。现代军事包括军事理论和军事科技两大部分。根据中华人民共和国教育部的定义，其下属的一级学科有：军事思想及军事历史、战略学、战役学、战术学、军队指挥学、军制学、军队政治工作、军事后勤学、军事装备学、军事训练学."}'

raw_str ='{"title": "化学键是什么", "text": "化学是一门研究物质的性质、组成、结构、以及变化规律的基础自然科学。化学研究的对象涉及物质之间的相互关系，或物质和能量之间的关联。传统的化学常常都是关于两种物质接触、变化，即化学反应，又或者是一种物质变成另一种物质的过程。这些变化有时会需要使用电磁波，当中电磁波负责激发化学作用。不过有时化学都不一定要关于物质之间的反应。光谱学研究物质与光之间的关系，而这些关系并不涉及化学反应。准确的说，化学的研究范围是包括分子、离子、原子、原子团在内的核-电子体系。「化学」一词，若单从字面解释就是「变化的学问」之意。化学主要研究的是化学物质互相作用的科学。化学如同物理皆为自然科学之基础科学。很多人称化学为「中心科学」，因为化学为部分科学学门的核心，连接物理概念及其他科学，如材料科学、纳米技术、生物化学等。研究化学的学者称为化学家。在化学家的概念中一切物质都是由原子或比原子更细小的物质组成，如电子、中子和质子。但化学反应都是以原子或原子团为最小结构进行的。若干原子通过某种方式结合起来可构成更复杂的结构，例如分子、离子或者晶体。当代的化学已发展出许多不同的学门，通常每一位化学家只专精于其中一、两门。在中学课程中的化学，化学家称为普通化学（，，）。普通化学是化学的导论。普通化学课程提供初学者入门简单的概念，相较于专业学门领域而言，并不甚深入和精确，但普通化学提供化学家直观、图像化的思维方式。即使是专业化学家，仍用这些简单概念来解释和思考一些复杂的知识。英语中的「化学」（）一字的语源有多种说法。一种说法认为是由「炼金术」（）得名的。英语中「」一词源于古法语的「」和阿拉伯语的「」，意为「形态变化的学问」（）。阿拉伯语中的「」一字则源于希腊语。亦有另一种说法认为英语中的「」一字源自埃及语中的「」，意思是「土」（）。在中国，「化学」一词最早出现在1857年墨海书馆出版的期刊《六合丛谈》。伟烈亚力提及王韬在其日记中记载了从戴德生处听闻的「化学」一词。一般认为中文中的“化学”一词是徐寿翻译英国人的书《化学鉴原》一书时发明的。「化学」一词被介绍到日本，取代了原先日语中的译法「舍密」。最早的化学要算是人类对火的研究。对于当时的人来说，火可以将一种物体变成另一种物体，所以成为了当时人最有兴趣研究的现象。如果没有火，人类不会发现到铁和玻璃的制造方法。人类发现了黄金这种贵重的金属之后，很多人转移研究怎样把其他物质变成黄金。公元前300年至1500年，炼金术士皆研究如何将一些便宜的金属转化成黄金，因此累积了金属的提取和处理有关的观察和技术。有些炼金术士主要的工作是制造药物，中国当时亦有所谓炼丹术。2000年前，人类已广泛使用金、银、汞、铜、铁和青铜。当时的人类文明，对于陶瓷、染色、酿造、造纸、火药等在工艺方面已有一定成就，在技术经验上，对物质变化的理解已有一定观察和文献累积。早期化学家收集了很多不同物质的资料。在17世纪以前，化学成就并不大（燃素说、炼金术），其中较有成就者如罗伯特·波义耳。到了1750年，化学仍带有神秘色彩，并为不正确的理论支配著。直到1773年，安托万-洛朗·德·拉瓦锡提出了质量守恒定律，并以氧化还原反应解释燃烧现象，推翻了盛行于中世纪的燃素说，才开启了现代化学之路；他因此被尊崇为「化学之父」。接著道尔顿整合当时的化学知识，并以自身的实验所得提出了划时代的原子说。此后，一些化学家相继发现了各种化学元素，后来门得列夫建立了元素周期表令化学视界更臻完备。1901年，化学家诺贝尔以其遗产成立了诺贝尔化学奖，以表扬在化学领域对人类有贡献者。现代化学始于20世纪初期蓬勃发展的量子力学。莱纳斯·鲍林引进量子力学解释化学键的本质，得以用波函数的线性叠加来描述。质子、中子和电子的发现，使化学真正由原子尺度来理解化学反应。量子力学和电子学的发展，使得许多新型仪器得以开发，来探索和分析化合物的结构和成分，如光谱仪、色谱仪、核磁共振仪和质谱仪等。当代化学大致分为四大学门，各学门又有许多延伸的子学门和应用化学领域。四大学门主要为：其他延展和应用的学门：一粒原子是由原子核及外围带负电荷的电子（称为核外电子）组成的粒子，一般而言是化学研究的最小尺度范畴。原子核通常是由质子和中子组成。与通常的物理概念不同的是，单一的质子在化学领域被认为是H（氕)原子核，也就是说原子核内必然含有质子，但可能不含中子。电子带负电荷，质子带正电荷，个数相同使得电荷平衡，令整个原子呈电中性。当核外电子数与原子核内质子数不相同时，则形成离子。通常认为离子也是原子的一种。拥有相同质子数的同一类原子被称为“元素”。例如，氢这种元素中所有原子都是只有一粒质子。这个概念换过来说亦可：所有原子核中有六粒质子的原子都是碳，所有原子核中有九十二粒质子的都是铀。元素亦有另一定义，就是所有不可以用化学方法分解的物质都是元素。在这么多种列举元素的方法中，最常用和最方便的莫过于元素周期表。周期表根据原子序数来排列原子，而原子序数就是一粒原子中质子的数量。因为这个奇怪的排列，排在一起的元素，无论是同一个直行、同一个横行还是纯粹在附近，都有一些大致上固定的关系。同一种元素可能有很多个不同的核素。它们的质子数相同而中子数不同，因而化学性质相同。但由于它们的中子数不同，造成原子核稳定性不同，而造成某些核素具备放射性。同一种元素的不同核素在元素周期表内占据同一个位置，因此同一种元素的不同核素互称同位素。例如H（氕）与H（氘）互称同位素。化学物质是指一种物体，它既确定了其化学组成，也确定了它的化学性质。严格的来讲，混合的化合物，元素等都不能算是化学物质，只能说是化学药品或者说化学制品。大多数我们日常生活碰到的化学品都是混合物，比如空气、合金、生物制品。物质的命名法在化学语言当中是最严格的一环。早在很久以前，化合物的命名是由其发现者自行决定的，这样则导致了命名的困难和混乱。而现在我们最常用的还是国际纯粹与应用化学联合会 () ()命名方法。它用一个命名系统让所有的化合物都有一个独有的名称和代码。有机化合物通过有机命名系统命名；而无机化合物通过无机命名系统命名。而通过化学索引服务（），我们可以轻松的通过CAS号()来找到每一个化合物的性质、特性、命名和结构。一个分子是化合物的最基本结构，不用化学方法是拆不开的。大部分分子都是由两个或以上原子组成，但是都有些特例，例如氦气分子，只有一个原子。这些原子，如果多于一个，是通过化学键结合。离子是带电荷的物质，可以由原子或分子失去或得到电子形成。正离子（例如钠离子Na）和负离子（例如氯离子Cl）结合可以成为电荷中性的盐（例如食盐NaCl）。有些离子是由几个原子组成，而它们进行化学作用的时候又不会分离，例如磷酸根离子（PO）、铵离子（NH）。气相的离子通常被称为等离子体。物质可以被分类为一种酸或者是一种碱。通常我们有几种进行酸碱分类定义的理论。其中最简单的要数阿累尼乌斯理论（），它认为：酸是能够在水当中电离出水合氢离子的物质；而相反碱则是在水当中电离出氢氧根离子的物质。而酸碱质子理论（）则认为酸是能够在化学反应中给其他物质氢离子的物质；而碱则是相应能得到氢离子的物质。第三种理论被称作是路易斯酸碱理论（)，它是基于形成化学键之上的。路易斯理论认为：酸是在键的形成当中接受了一对电子；而碱则是在形成键的过程中给予了其他物质一对电子。因此，一个物质如果对于不同的酸碱理论来说，可能在此是酸，在另外一个理论来说却是碱。酸性强度的衡量方法主要有两种：第一种是阿累尼乌斯定义的也就是我们最常用的pH，它是通过衡量一个溶液当中氢离子的浓度来确定酸性的大小。它的计算方法是pH=-log[H]，也就是pH等于氢离子浓度的负对数（以10为底）。因此可以说，拥有更高浓度的氢离子溶液，其pH越低而酸性更强。第二种是定义，也就是酸解离常数（K），它衡量的是物质作为酸的时候给予氢离子的能力。因此一个酸性越强的物质，其K更高，更具有给予氢离子的的倾向。同样的我们可以用pOH代替pH, K代替K来说明碱性强度。氧化还原的概念和一个物质的原子获取或者给予电子的能力有关。物质拥有氧化其他物质的能力就被成为氧化性，而此物质被成为氧化剂（)，或者成为氧化物。一个氧化剂能够将电子从其他的物质上移走。相应的，具有还原其他物质的物质被称作有还原性而成为还原剂（）或者成为还原物。一个还原试剂能够传递给其他物质电子并且氧化自身。而正因为其“给予”了其他物质电子，它还被称为供电子物。氧化还原的性质与氧化数（）有关－－其实真正的给予或者获取完成的电子并不存在。所以，氧化过程被定义为增加了氧化数，而还原则是降低的氧化数。化学品泛指一切有确实化学构造及化学成份的物质，所以又称化学物质。它们可以是元素、化合物或混合物。日常生活中，我们会遇到的东西多数都是混合物，例如合金。化合物是一些以不同元素用固定比例结合而成的物质。成份的比例决定了它的化学特性。例如水是用氢同氧以二比一组合而成，组成水分子的三个原子之间构成了104.5度的健角。不同化合物及元素之间的变化称为化学反应。摩尔（，仅台湾地区使用莫耳一词）是物质的量的国际单位，符号为mol。1摩尔是所含基本微粒个数与12克的碳-12(formula_1)中所含原子个数相等的一系统物质的量。使用摩尔时，应指明基本微粒，可以是分子、原子、离子、电子或其他基本微粒，也可以是基本微粒的特定组合体。1摩尔物质中所含基本微粒的个数等于阿伏伽德罗常数，符号为N，数值是6.0214129×10，常取6.02×10。一种物质的摩尔质量与分子量，在使用国际单位制时，在数值上相等。化学键是指组成分子或材料的粒子之间互相作用的力量，其中粒子可以是原子、离子或是分子。化学键的物理本质来自于粒子和粒子之间的静电力，量子力学上意指原子间电子的波函数线性叠加。化学键是化学最重要的概念之一，物理理论本质由莱纳斯·鲍林建立。化学家为能简洁表述化学键并规避量子力学的复杂性，将化学键分类为共价键、离子键和金属键，较弱的键结如氢键及较特殊的配位相互作用等。无论分类为何，其物理本质都是相同的。分子间力是不同分子之间的作用力，主要有氢键，范德华力，亲水作用／疏水作用等，这种作用力比化学键弱，容易打开或重新组合，但是是形成分子空间排列和架构的重要作用力，是现代化学的重要研究方向之一。物质有时会是液体，有时会是固体，有时会是气体，这些叫作物质的相态。一件物质是否软、透不透光、透光的话它的折射率是多少，这些都是一件物质的物理特性。总而言之，物理特性即是一种物质不靠化学作用都可以断定到的特性。化学反应，亦称化学变化是一种物质转变为另一种物质的过程，涉及分子中原子的交换和化学键的转移、形成或消失。化学反应形成的改变既可令很多独立的分子结合，也可将一个较大型的分子拆开成为很多独立的小分子，甚至是同一分子内有原子移动，即使原子的数量没有改变，但仍会构成化学反应。虽然平衡概念在科学各领域都得到广泛的应用，但在化学中，化学平衡是指化学成分中出现多种不同状态的可能性，例如在可以彼此反应的几种化合物的混合物中，或当物质可以以多于一种相态存在的时候。即使有著不变的化学组成，但在平衡系统中的化学物质通常并非处于静止状态；这些物质的分子会互相继续进行反应，从而产生动态平衡。因此，化学平衡描述了诸如化学成分之类的参数随时间保持不变的状态。化学反应的守恒必须符合物理守恒定律，反应前后应符合：化学工业（化工）是当代经济活动当中重要的一部分。全球50大化学品制造商在2004年共销售了5870亿美元的业绩，其中利润占据了8.1%，其中研发成本占据了2.1% 其他还有诸如放射分析化学、同位素化学、辐射化学、核燃料、反应堆和裂变产物化学、地球化学、海洋化学、大气化学、环境化学、宇宙化学、星际化学、药物化学、神经化学、农业化学、石油化学、木材化学、土壤化学、煤化学、食品化学、化学地理学、天体化学、岩石化学、空间化学及胶体与界面化学等"}'

test_data = json.loads(raw_str)

context = test_data["text"]
question = test_data["title"]

max_length = 384
stride = 128

batch_data = tokenizer(
        question,
        context,
        max_length=max_length,
        truncation="only_second",
        stride=stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length", 
        return_tensors="pt"
)
print(batch_data.keys())
offset_mapping = batch_data.pop('offset_mapping').numpy().tolist()
sample_mapping = batch_data.pop('overflow_to_sample_mapping')
print(batch_data.keys(), len(offset_mapping))

offsets=[]

for i in range(len(batch_data['input_ids'])):
        sequence_ids = batch_data.sequence_ids(i)
        offset = offset_mapping[i]
        offset_mapping[i] = [
                o if sequence_ids[k] == 1 else None for k, o in enumerate(offset)
        ]
        offsets += offset_mapping[i]

pred_start_logits, pred_end_logit = model(batch_data)
print(pred_start_logits.shape, pred_end_logit.shape)

n_best = 20
max_answer_length = 30
answers=[]
start_logits_all=[]
end_logits_all=[]

with torch.no_grad():
        start_logits_all=pred_start_logits.cpu().numpy()
        end_logits_all=pred_end_logit.cpu().numpy()
        #print(len(offsets))
        start_logit = np.concatenate(start_logits_all)
        end_logit = np.concatenate(end_logits_all)
        start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist()
        end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist()
        for start_index in start_indexes:
                for end_index in end_indexes:
                        if offsets[start_index] is None or offsets[end_index] is None:
                                continue
                        if (end_index < start_index or end_index - start_index + 1 > max_answer_length):
                                continue
                        answers.append({
                                "start": offsets[start_index][0], 
                                "text": context[offsets[start_index][0] : offsets[end_index][1]], 
                                "logit_score": start_logit[start_index] + end_logit[end_index],
                                })
        best_answer = max(answers, key=lambda x: x["logit_score"])
        print(answers)
        print(best_answer)


